"""
合同语义审查
"""
import json
from typing import Any, Dict
from core.llm_answer import answer_with_llm
from core.milvus_vector import MilvusVector
from core.config import settings
from core.reranker import ReRanker
from loguru import logger
import rapidfuzz
import re


def check_prompt(contract_text, rule_name, extract_hint, rule_hint) -> str:

    if extract_hint != '':
        extract_text = f'按照如下指示：`{extract_hint}`，并根据合同内容排版的特点，准确找出`{rule_name}`所意指的文本内容'
    else:
        extract_text = f'根据合同内容排版的特点，准确找出`{rule_name}`所意指的文本内容'

    if rule_name.find('合同名称') != -1:
        extract_text = extract_text + f'，协议也是合同的一种，通常，`{rule_name}`位于合同的开始部分'

    if re.search('人[^\u4E00-\u9F5A]*$', rule_name) is not None:
        extract_text += f'，注意：所找出的文本内容必须包含有效的人名'

    prompt = f'''
      忘掉之前的所有设定和对话，从现在开始，你是一名专业的合同审查助手。现合同中有以下条款段落(按顺序排列)：
      `{contract_text}`
      
      你需要基于你的专业知识，{extract_text}，并判断这些文本内容是否符合如下要求：
      `{rule_hint}`
      
      结果以 JSON 格式输出，包含`content`、`approved` 和 `reason` 三个字段，字段说明如下：
      - content：`{rule_name}`所指向的文本内容，如果未找到或所指向的内容未填写，则为空字符串。
      - approved：值为 success 表示符合审查规则，或 fail 表示不符合审查规则。
      - reason：简要说明判断的理由或相关描述内容。
      '''
    return prompt


async def rule_review_api(fileId: str, rule: Dict[str, Any]) -> Dict:
    """
    合同要素的语义审查
    """
    info = {'fileId': fileId, 'rule': rule}
    logger.info(f'要素语义审查: {json.dumps(info, ensure_ascii=False)}')

    # 召回要素相关的片段
    milvus = MilvusVector()
    rule_name = rule['name']
    extract_hint = rule.get('extractHint', '') or ''
    check_hint = rule.get('checkHint', '') or ''

    query_parts = []
    if extract_hint != '':
        query_parts.append(f'按照如下指示：`{extract_hint}`')
    query_parts.append(f'找出`{rule_name}`相关的文本内容')
    if check_hint != '':
        query_parts.append(f'并要求：{check_hint}')
    query = '，'.join(query_parts) + "。"

    documents = milvus.hybrid_search(fileId, query, top_k=30)

    # 文档重排
    reranker = ReRanker()
    documents = reranker.rerank(query, documents, settings.TOP_K)

    if settings.DEBUG:
        logger.debug('召回的文本片段：')
        for doc in documents:
            logger.debug(doc)

    # 调用大模型进行要素抽取和语义审查
    prompt = None
    related_dict = None
    for k in reversed(range(0, len(documents))):
        related_texts, related_dict = reorganize(documents[0: k])
        prompt = check_prompt('\n'.join(related_texts), rule_name, extract_hint, check_hint)
        if len(prompt) < settings.LLM_MAX_INPUT:
            break

    answer = await answer_with_llm(prompt, parse_type='dict')
    if settings.DEBUG:
        logger.debug(f'prompt：{prompt}')
        logger.debug(f'answer：{answer}')

    # 与要素最相关的内容
    content = answer.get('content', '')
    if content != '' and len(documents) != 0:
        partial_ratios = dict()
        key2id = dict()
        for key, block in related_dict.items():
            partial_ratios[key] = rapidfuzz.fuzz.partial_ratio(block['content'], content)
            key2id[key] = block['id']
        sored_ratios = sorted(partial_ratios.items(), key=lambda x: -x[1])
        paragraph_id = key2id[sored_ratios[0][0]]
        sorted_ids = list(dict.fromkeys([item[0] for item in sored_ratios]))
    else:
        paragraph_id = None
        sorted_ids = list(dict.fromkeys([doc.metadata['id'] for doc in documents]))

    if check_hint != '':
        reg = re.search('(success|fail)', str(answer), re.IGNORECASE)
        check_label = reg.group(1) if reg is not None else None
        reason = answer.get('reason', str(answer))
        if 'reason' not in answer:
            error_info = { 'prompt': prompt, 'answer': answer}
            logger.error(f"大模型回答错误：{error_info}")
    else:
        check_label = None
        reason = None

    # 返回结果
    check_result = {
        'fileId': fileId,             # 合同文件id
        'id': rule.get('id', None),   # 要素id
        'name': rule_name,            # 要素名称
        'extractHint': extract_hint,  # 要素抽取的hint
        'checkHint': check_hint,      # 语义审查的hint
        'related_ids': sorted_ids,    # 召回的相关段落id
        'content': re.sub('###?\s+', '', content),           # 要素的内容
        'paragraph_id': paragraph_id, # 要素的内容关联的段落id
        'check_label': check_label,   # 判定结果
        'reason': reason,             # 判定理由
    }
    logger.info(f'语义审查结果: {json.dumps(check_result, ensure_ascii=False)}')
    return check_result


def reorganize(documents) :
    # 按合同原始顺序拼接成markdown形式的文本, 标题也需整合进去
    retrieved = dict()
    related_dict = dict()
    for doc in documents:
        related_blocks = doc.metadata['related_blocks']
        for key, block in related_blocks.items():
            start_index = int(key.split('-')[1])
            content = block['content']

            block_index = int(key.split('-')[0])
            if block_index not in retrieved:
                retrieved[block_index] = dict()
            retrieved[block_index][start_index] = content
            related_dict[key] = block

    retrieved = sorted(retrieved.items(), key=lambda x: x[0])
    related_texts = []
    for block_index, content_items in retrieved:
        content_items = sorted(content_items.items(), key=lambda x: x[0])
        related_texts.append(''.join([content for _, content in content_items]))

    return related_texts, related_dict